import pandas as pd
import numpy as np
import ta
import pandas_ta as ta_lib

# # 假设你已经有了一个包含市场数据的DataFrame，这里我们创建一个示例DataFrame
# dates = pd.date_range("20240101", periods=100)
# data = pd.DataFrame(index=dates)
# data["Close"] = np.random.randn(100).cumsum() + 100

import yfinance as yf

# 获取历史数据（示例用AAPL）
data = yf.download("AAPL", start="2025-01-01", end="2025-02-26")

# 计算MACD指标
data["MACD"], data["MACD_signal"], _ = ta.trend.macd(data["Close"])

# 计算KDJ指标（注意：pandas_ta中的KDJ指标可能与你所知的传统KDJ有所不同）
kdj = ta_lib.KDJ(data["Close"], k=9, d=3, j=3)
data["K"] = kdj.kdj_k
data["D"] = kdj.kdj_d
data["J"] = kdj.kdj_j

# 计算RSI指标
data["RSI"] = ta.momentum.rsi(data["Close"], window=14)


# 定义交易策略
def trading_strategy(data):
    signals = pd.DataFrame(index=data.index)
    signals["signal"] = 0

    # MACD金叉买入，死叉卖出
    signals["signal"][1:] = np.where(data["MACD"][1:] > data["MACD_signal"][1:], 1, 0)
    signals["positions"] = signals["signal"].diff()

    # 结合KDJ指标：J值大于D值且RSI小于70时买入；J值小于D值且RSI大于30时卖出
    buy_condition = (data["J"] > data["D"]) & (data["RSI"] < 70)
    sell_condition = (data["J"] < data["D"]) & (data["RSI"] > 30)

    signals.loc[buy_condition, "signal"] = 1
    signals.loc[sell_condition, "signal"] = -1

    return signals


# 生成交易信号
signals = trading_strategy(data)

# 打印交易信号
print(signals)

# 你可以进一步根据这些信号来执行交易，比如使用回测框架来评估策略的表现
